#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import time
import sys
import os
import matplotlib.pyplot as plt
import matplotlib.path as mpath
import matplotlib.patches as mpatches
from datetime import datetime
import logging

# 设置日志
log_filename = f"robot_path_planning_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
logging.basicConfig(
    filename=log_filename,
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

# 创建控制台处理器，同时输出到控制台和文件
console = logging.StreamHandler()
console.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)

# 设置Matplotlib字体和样式
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用于显示中文
plt.rcParams['axes.unicode_minus'] = False    # 解决保存图像是负号'-'显示为方块的问题
try:
    # 设置指定的字体路径
    font_path = r"E:\anaconda3\Lib\site-packages\qtawesome\fonts\elusiveicons-webfont-2.0.ttf"
    if os.path.exists(font_path):
        from matplotlib.font_manager import FontProperties
        custom_font = FontProperties(fname=font_path)
        logging.info(f"已加载自定义字体: {font_path}")
    else:
        logging.warning(f"指定的字体文件不存在: {font_path}")
except Exception as e:
    logging.warning(f"设置自定义字体时出错: {e}")

# 创建用于可视化的钻石标记
diamond_marker = mpath.Path([(-0.5, 0), (0, 0.5), (0.5, 0), (0, -0.5), (-0.5, 0)], closed=True)

# 简化导入，避免使用模块
# 原来的导入
# from robot_arm.arm import RobotArm
# from path_planning.planner import AStarPlanner
# from optimization.optimizer import OptimizationSystem

# 创建简单的替代类
class RobotArm:
    def __init__(self):
        logging.info("初始化机械臂")
        self.arm_length = [400, 350, 300]  # 三个关节的长度(mm)
        self.dof = 3  # 自由度
    
    def calculate_joint_positions(self, angles):
        """计算各关节位置"""
        positions = []
        # 第一个关节固定在底座
        pos = np.array([0, 0, 0])
        positions.append(pos.copy())
        
        current_angle = 0
        for i in range(self.dof):
            current_angle += angles[i]
            dx = self.arm_length[i] * np.cos(current_angle)
            dy = self.arm_length[i] * np.sin(current_angle)
            pos[0] += dx
            pos[1] += dy
            positions.append(pos.copy())
        
        return positions

class AStarPlanner:
    def __init__(self, map_width=30, map_height=30):
        self.map_width = map_width
        self.map_height = map_height
        self.obstacles = set()
        logging.info(f"初始化路径规划器，地图大小: {map_width}x{map_height}")
    
    def add_obstacle(self, x, y):
        self.obstacles.add((x, y))
    
    def plan_path(self, start, goal):
        logging.info(f"规划路径: {start} -> {goal}")
        # 简单实现：这里使用A*算法,但简化为直接返回起点和终点
        # 真实项目中应实现完整的A*算法
        return [start, goal]

class OptimizationSystem:
    def __init__(self, robot):
        self.robot = robot
        self.obstacles = []
        logging.info("初始化优化系统")
    
    def add_obstacle(self, pos):
        self.obstacles.append(pos)
    
    def optimize_path(self, path):
        logging.info(f"优化路径，路径点数: {len(path)}")
        # 简单实现：返回原路径和示例误差与能耗
        return path, 0.5, 1.2

def visualize_path(obstacles, start_pos, targets, all_paths, filename="path_visualization.png"):
    """可视化路径规划结果"""
    logging.info(f"开始生成可视化图像: {filename}")
    
    plt.figure(figsize=(12, 10))
    
    # 绘制网格背景
    grid_size = 100  # 网格大小(mm)
    
    # 计算坐标范围
    max_x = 2000
    max_y = 2000
    
    if all_paths and len(all_paths) > 0:
        max_x = max(max_x, max([p[0] for p in all_paths]) + 500)
        max_y = max(max_y, max([p[1] for p in all_paths]) + 500)
    
    if isinstance(targets, np.ndarray) and targets.size > 0:
        max_x = max(max_x, targets[:, 0].max() + 500)
        max_y = max(max_y, targets[:, 1].max() + 500)
    
    # 绘制网格线
    for x in range(0, int(max_x) + grid_size, grid_size):
        plt.axvline(x, color='lightgray', linestyle='-', linewidth=0.5)
    for y in range(0, int(max_y) + grid_size, grid_size):
        plt.axhline(y, color='lightgray', linestyle='-', linewidth=0.5)
    
    # 添加网格坐标标签
    for i in range(0, int(max_x) // grid_size + 1):
        for j in range(0, int(max_y) // grid_size + 1):
            plt.text(i * grid_size + grid_size/2, j * grid_size + grid_size/2, 
                     f"({i},{j})", ha='center', va='center', fontsize=8, color='gray')
    
    # 绘制障碍物
    if isinstance(obstacles, np.ndarray) and obstacles.size > 0:
        obstacle_x = obstacles[:, 0]
        obstacle_y = obstacles[:, 1]
        plt.scatter(obstacle_x, obstacle_y, color='red', marker='s', s=80, label='障碍物')
        logging.info(f"绘制 {len(obstacles)} 个障碍物")
    
    # 绘制起点
    plt.scatter(start_pos[0]*100, start_pos[1]*100, color='green', marker='o', s=160, label='起点')
    logging.info(f"绘制起点: ({start_pos[0]*100}, {start_pos[1]*100})")
    
    # 绘制目标点
    if isinstance(targets, np.ndarray) and targets.size > 0:
        for i in range(len(targets)):
            plt.scatter(targets[i][0], targets[i][1], color='blue', marker=diamond_marker, s=140, label=f'目标点' if i==0 else "")
            plt.annotate(f'T{i+1}', (targets[i][0], targets[i][1]), fontsize=12)
            logging.info(f"绘制目标点{i+1}: ({targets[i][0]}, {targets[i][1]})")
    
    # 绘制路径
    if all_paths and len(all_paths) > 0:
        path_x = [point[0] for point in all_paths]
        path_y = [point[1] for point in all_paths]
        
        # 将路径分段，用不同颜色表示
        segments = []
        current_segment = []
        
        for i in range(len(all_paths)-1):
            current_segment.append((path_x[i], path_y[i]))
            
            # 检测是否到达目标点（简化判断：如果下一个点接近任一目标点，则认为是新段开始）
            is_target = False
            if isinstance(targets, np.ndarray) and targets.size > 0:
                for t in targets:
                    if (abs(path_x[i+1] - t[0]) < 50 and 
                        abs(path_y[i+1] - t[1]) < 50):
                        is_target = True
                        break
            
            if is_target and current_segment:
                current_segment.append((path_x[i+1], path_y[i+1]))
                segments.append(current_segment)
                current_segment = []
            
        # 添加最后一段
        if current_segment:
            if len(all_paths) > len(current_segment):
                current_segment.append((path_x[-1], path_y[-1]))
            segments.append(current_segment)
        
        # 如果没有分段，则整个路径作为一个段
        if not segments and all_paths:
            segments = [[(p[0], p[1]) for p in all_paths]]
        
        # 绘制分段路径
        colors = ['blue', 'purple', 'orange', 'cyan', 'magenta', 'yellow']
        
        for i, segment in enumerate(segments):
            if len(segment) < 2:
                continue
                
            color = colors[i % len(colors)]
            x_vals = [point[0] for point in segment]
            y_vals = [point[1] for point in segment]
            plt.plot(x_vals, y_vals, color=color, linewidth=2, 
                    label=f'路径段{i+1}' if i==0 else "")
            
            # 添加箭头表示方向
            if len(segment) > 1:
                mid_idx = len(segment) // 2
                dx = segment[mid_idx][0] - segment[mid_idx-1][0]
                dy = segment[mid_idx][1] - segment[mid_idx-1][1]
                plt.arrow(segment[mid_idx-1][0], segment[mid_idx-1][1], dx, dy, 
                         head_width=50, head_length=50, fc=color, ec=color)
        
        logging.info(f"绘制 {len(segments)} 段路径，总路径点数: {len(all_paths)}")
    
    plt.title('机械臂路径规划可视化 (栅格大小: 100mm×100mm)')
    plt.xlabel('X坐标 (mm)')
    plt.ylabel('Y坐标 (mm)')
    plt.grid(True)
    plt.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=3)
    
    # 添加说明文本
    plt.figtext(0.5, 0.01, 
               "注: 0表示没有障碍物(可行区域)，1表示有障碍物，Start表示机器臂初始位置，End表示目标点位置，每格的大小为100mm×100mm",
               ha='center', fontsize=9)
    
    # 保存图像
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    logging.info(f"可视化图像已保存: {filename}")
    
    # 显示图像
    plt.close()

def solve_problem4(excel_file_path, sheet_name="Sheet2"):
    """
    解决问题4：识别Excel表格中的障碍物、起点和目标点，规划最优路径
    
    参数:
        excel_file_path: Excel文件路径
        sheet_name: 工作表名称
    
    返回:
        total_error: 总误差
        total_energy: 总能耗
        all_paths: 所有路径点
    """
    try:
        logging.info(f"尝试读取Excel文件: {excel_file_path}, 工作表: {sheet_name}")
        
        # 读取Excel文件的指定工作表
        df = pd.read_excel(excel_file_path, sheet_name=sheet_name, header=None)
        
        if df.empty:
            logging.error("Excel 文件为空")
            raise Exception("Excel 文件为空")
        
        logging.info(f"成功读取Excel文件，表格大小: {df.shape}")
        
        # 初始化数据结构
        obstacle_pos = []
        targets = []
        start_pos = None
        
        # 遍历Excel表格，识别障碍物、起点和目标点
        for row_idx in range(df.shape[0]):
            for col_idx in range(df.shape[1]):
                cell_value = df.iloc[row_idx, col_idx]
                
                # 跳过空单元格
                if pd.isna(cell_value):
                    continue
                
                # 障碍物标记为1
                if cell_value == 1:
                    # 转换为毫米坐标
                    x = col_idx * 100
                    y = row_idx * 100
                    obstacle_pos.append([x, y, 0])
                
                # 起点标记为"Start"
                elif isinstance(cell_value, str) and cell_value.lower() == "start":
                    # 转换为网格坐标
                    start_pos = (col_idx, row_idx)
                    logging.info(f"找到起点: {start_pos}")
                
                # 目标点标记为"target1", "target2"等
                elif isinstance(cell_value, str) and cell_value.lower().startswith("target"):
                    # 转换为毫米坐标，z坐标设为200mm
                    x = col_idx * 100
                    y = row_idx * 100
                    z = 200
                    targets.append([x, y, z])
                    logging.info(f"找到目标点: ({x}, {y}, {z})")
        
        # 将列表转换为NumPy数组
        obstacle_pos = np.array(obstacle_pos)
        
        # 检查是否找到了起点和至少一个目标点
        if start_pos is None:
            logging.warning("表格中没有找到起点，使用默认值 (0, 0)")
            start_pos = (0, 0)
        
        if len(targets) == 0:
            logging.warning("表格中没有找到目标点，使用默认值")
        targets = np.array([
            [1500, 1200, 200],
                [1000, 800, 200],
                [800, 1500, 200]
            ])
        else:
            targets = np.array(targets)
        
        logging.info(f"找到 {len(obstacle_pos)} 个障碍物")
        logging.info(f"找到 {len(targets)} 个目标点")
        
        # 创建机械臂实例
        robot = RobotArm()
        
        # 创建路径规划器
        planner = AStarPlanner(map_width=max(30, df.shape[1]), map_height=max(30, df.shape[0]))
        
        # 创建优化系统
        optimizer = OptimizationSystem(robot)
        
        # 添加障碍物到路径规划器
        if len(obstacle_pos) > 0:
        for obs in obstacle_pos:
            planner.add_obstacle(int(obs[0]/100), int(obs[1]/100))
                optimizer.add_obstacle(np.array([obs[0], obs[1]]))
            logging.info(f"已添加 {len(obstacle_pos)} 个障碍物到规划器")
        else:
            logging.warning("没有添加任何障碍物，路径规划可能不准确")
        
        # 将起始位置转换为毫米坐标
        current_pos = np.array([start_pos[0] * 100, start_pos[1] * 100, 0])
        logging.info(f"机械臂初始位置: {current_pos}")
        
        # 初始化路径和计数器
        all_paths = []
        total_error = 0
        total_energy = 0
        path_counter = 0
        
        # 遍历所有目标点规划路径
        for i, target in enumerate(targets):
            logging.info(f"\n规划到目标点 {i+1}: {target}")
            
            try:
                # 将当前位置和目标位置转换为栅格坐标
                start_grid = (int(current_pos[0]/100), int(current_pos[1]/100))
                goal_grid = (int(target[0]/100), int(target[1]/100))
                
                # 规划路径
                path = planner.plan_path(start_grid, goal_grid)
            
            if not path:
                    logging.warning(f"无法规划从 {start_grid} 到 {goal_grid} 的路径，尝试替代方案")
                    # 替代方案：直接连接当前位置和目标位置
                    path = [start_grid, goal_grid]
                
                # 将网格路径转换为毫米坐标
                mm_path = []
                for point in path:
                    # 转换XY坐标，对应的Z坐标采用线性插值
                    x = point[0] * 100
                    y = point[1] * 100
                    
                    # 计算Z坐标：起点Z=0，终点Z=target[2]
                    progress = len(mm_path) / max(1, len(path) - 1)
                    z = target[2] * progress
                    
                    mm_path.append([x, y, z])
                
                # 优化路径
                try:
                    optimized_path, error, energy = optimizer.optimize_path(mm_path)
                    logging.info(f"路径优化成功: 误差={error:.2f}, 能耗={energy:.2f}")
                    
                    # 累计误差和能耗
                    total_error += error
                    total_energy += energy
                    path_counter += 1
                    
                    # 保存优化后的路径
                    all_paths.extend(optimized_path)
                    
                    # 更新当前位置为路径的终点
                    current_pos = optimized_path[-1]
                    
                except Exception as opt_e:
                    logging.error(f"路径优化失败，使用原始路径: {opt_e}")
                    all_paths.extend(mm_path)
                    current_pos = mm_path[-1]
            
            except Exception as path_e:
                logging.error(f"路径规划失败: {path_e}")
                # 错误处理：不再递归调用，而是跳过当前目标点
                logging.warning(f"跳过目标点 {i+1}")
                continue
        
        # 规划返回起点的路径
        try:
            logging.info("\n规划返回起点的路径")
            start_grid = (int(current_pos[0]/100), int(current_pos[1]/100))
            goal_grid = start_pos
            
            # 规划返回路径
            return_path = planner.plan_path(start_grid, goal_grid)
            
            if not return_path:
                logging.warning(f"无法规划返回起点的路径，使用直线")
                return_path = [start_grid, goal_grid]
            
            # 转换为毫米坐标
            mm_return_path = []
            for point in return_path:
                # 返回时Z坐标从当前高度逐渐降至0
                progress = len(mm_return_path) / max(1, len(return_path) - 1)
                start_z = current_pos[2]
                end_z = 0
                z = start_z * (1 - progress) + end_z * progress
                
                mm_return_path.append([point[0] * 100, point[1] * 100, z])
            
            # 优化返回路径
            try:
                optimized_return, error, energy = optimizer.optimize_path(mm_return_path)
                logging.info(f"返回路径优化成功: 误差={error:.2f}, 能耗={energy:.2f}")
                
                total_error += error
                total_energy += energy
                path_counter += 1
                
                all_paths.extend(optimized_return)
            
            except Exception as opt_e:
                logging.error(f"返回路径优化失败: {opt_e}")
                all_paths.extend(mm_return_path)
        
        except Exception as return_e:
            logging.error(f"规划返回路径失败: {return_e}")
            # 直接结束，不再尝试返回起点
        
        # 计算平均误差和能耗
        avg_error = total_error / max(1, path_counter)
        avg_energy = total_energy / max(1, path_counter)
        
        logging.info(f"\n规划完成:")
        logging.info(f"总路径点数: {len(all_paths)}")
        logging.info(f"总误差: {total_error:.2f}")
        logging.info(f"总能耗: {total_energy:.2f}")
        logging.info(f"平均误差/目标: {avg_error:.2f}")
        logging.info(f"平均能耗/目标: {avg_energy:.2f}")
        
        # 生成可视化
        visualize_path(obstacle_pos, start_pos, targets, all_paths, "problem4_path_visualization.png")
        
        # 可视化机械臂路径
        visualize_arm_path(robot, obstacle_pos, start_pos, targets, all_paths, "problem4_arm_path.png")
        
        return total_error, total_energy, all_paths
    
    except Exception as e:
        logging.error(f"处理Excel文件时发生错误: {e}")
        import traceback
        traceback.print_exc()
        logging.error(traceback.format_exc())
        
        # 创建默认数据
        logging.info("使用默认数据继续计算")
        
        obstacles = np.array([
            [300, 300, 0],
            [700, 500, 0],
            [500, 900, 0]
        ])
        
        targets = np.array([
            [1500, 1200, 200],
            [1000, 800, 200],
            [800, 1500, 200]
        ])
        
        # 创建机械臂实例
        robot = RobotArm()
        
        # 创建路径规划器
        planner = AStarPlanner(map_width=30, map_height=30)
        
        # 创建优化系统
        optimizer = OptimizationSystem(robot)
        
        # 添加障碍物
        for obs in obstacles:
            planner.add_obstacle(int(obs[0]/100), int(obs[1]/100))
            optimizer.add_obstacle(np.array([obs[0], obs[1]]))
        
        # 初始化数据
        current_pos = np.array([0, 0, 0])
        all_paths = []
        total_error = 0
        total_energy = 0
        path_counter = 0
        
        # 处理所有目标点
        for i, target in enumerate(targets):
            try:
                start_grid = (int(current_pos[0]/100), int(current_pos[1]/100))
                goal_grid = (int(target[0]/100), int(target[1]/100))
                
                path = planner.plan_path(start_grid, goal_grid)
                
                if not path:
                    path = [start_grid, goal_grid]
                
                mm_path = []
                for point in path:
                    progress = len(mm_path) / max(1, len(path) - 1)
                    z = target[2] * progress
                    mm_path.append([point[0] * 100, point[1] * 100, z])
                
                try:
                    optimized_path, error, energy = optimizer.optimize_path(mm_path)
                    total_error += error
                    total_energy += energy
                    path_counter += 1
                    all_paths.extend(optimized_path)
                    current_pos = optimized_path[-1]
                except:
                    all_paths.extend(mm_path)
                    current_pos = mm_path[-1]
            except:
                continue
        
        logging.info(f"\n使用默认数据计算结果:")
        logging.info(f"总路径点数: {len(all_paths)}")
        logging.info(f"总误差: {total_error:.2f}")
        logging.info(f"总能耗: {total_energy:.2f}")
        
        # 生成默认数据的可视化
        visualize_path(obstacles, (0, 0), targets, all_paths, "problem4_default_path.png")
        
        return total_error, total_energy, all_paths

def visualize_arm_path(robot, obstacles, start_pos, targets, all_paths, filename="problem4_arm_path.png"):
    """可视化机械臂在多个目标点的动作"""
    logging.info(f"开始生成机械臂路径可视化: {filename}")
    
    # 创建图像
    plt.figure(figsize=(15, 10))
    
    # 绘制栅格地图背景
    grid_size = 100  # 毫米
    max_x = 2000
    max_y = 2000
    
    if all_paths and len(all_paths) > 0:
        max_x = max(max_x, max([p[0] for p in all_paths]) + 200)
        max_y = max(max_y, max([p[1] for p in all_paths]) + 200)
    
    if isinstance(targets, np.ndarray) and targets.size > 0:
        max_x = max(max_x, targets[:, 0].max() + 200)
        max_y = max(max_y, targets[:, 1].max() + 200)
    
    # 绘制栅格
    for x in range(0, int(max_x) + grid_size, grid_size):
        plt.axvline(x, color='lightgray', linestyle='-', linewidth=0.5)
    for y in range(0, int(max_y) + grid_size, grid_size):
        plt.axhline(y, color='lightgray', linestyle='-', linewidth=0.5)
    
    # 绘制障碍物
    if isinstance(obstacles, np.ndarray) and obstacles.size > 0:
        for i in range(len(obstacles)):
            plt.scatter(obstacles[i][0], obstacles[i][1], color='red', marker='s', s=80, 
                       label='障碍物' if i == 0 else "")
    
    # 绘制起点
    plt.scatter(start_pos[0]*100, start_pos[1]*100, color='green', marker='o', s=160, label='起点(Start)')
    
    # 绘制目标点
    if isinstance(targets, np.ndarray) and targets.size > 0:
        for i in range(len(targets)):
            plt.scatter(targets[i][0], targets[i][1], color='blue', marker=diamond_marker, s=140, 
                       label=f'目标点(Target)' if i==0 else "")
            plt.annotate(f'T{i+1}', (targets[i][0], targets[i][1]), fontsize=12)
    
    # 绘制机械臂路径
    if all_paths and len(all_paths) > 0:
        # 简化：只绘制XY平面上的路径
        path_x = [p[0] for p in all_paths]
        path_y = [p[1] for p in all_paths]
        plt.plot(path_x, path_y, 'b-', linewidth=2, label='底座移动路径')
        
        # 添加方向箭头
        arrow_step = max(1, len(all_paths) // 10)  # 每隔几个点添加一个箭头
        for i in range(0, len(all_paths)-1, arrow_step):
            plt.arrow(path_x[i], path_y[i], 
                     path_x[i+1]-path_x[i], path_y[i+1]-path_y[i],
                     head_width=30, head_length=30, fc='blue', ec='blue')
    
    plt.title('机器人底座移动路径与货物抓取任务可视化 (栅格大小: 100mm×100mm)')
    plt.xlabel('X坐标 (mm)')
    plt.ylabel('Y坐标 (mm)')
    plt.grid(True)
    plt.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=3)
    
    # 添加说明文本
    plt.figtext(0.5, 0.01, 
               "注: 0表示没有障碍物(可行区域)，1表示有障碍物，Start表示机器臂初始位置，Target表示目标点位置，每格的大小为100mm×100mm",
               ha='center', fontsize=9)
    
    plt.savefig(filename, dpi=300, bbox_inches='tight')
    logging.info(f"机械臂路径可视化已保存: {filename}")
    
    plt.close()

def main():
    """主函数，用于测试和演示"""
    # 设置Excel文件路径（默认路径或从命令行参数获取）
    if len(sys.argv) > 1:
        excel_file_path = sys.argv[1]
    else:
        # 默认路径
        excel_file_path = "F:/kit/喜/A题附件.xlsx"
    
    # 确保文件存在
    if not os.path.exists(excel_file_path):
        logging.error(f"错误: 找不到Excel文件 '{excel_file_path}'")
        sys.exit(1)
    
    logging.info(f"读取Excel文件: {excel_file_path}")
    
    try:
        # 解决问题4
        error, energy, paths = solve_problem4(excel_file_path, sheet_name="Sheet2")
        
        # 输出结果
        logging.info("\n最终结果:")
        logging.info(f"总误差: {error:.2f}")
        logging.info(f"总能耗: {energy:.2f}")
        
        # 生成摘要信息
        print("\n" + "="*50)
        print("问题4求解结果摘要")
        print("="*50)
        print(f"总末端误差: {error:.2f}")
        print(f"总能耗: {energy:.2f}")
        print(f"总路径点数: {len(paths)}")
        print("="*50)
        print("已生成可视化结果:")
        print("1. problem4_path_visualization.png - 显示栅格路径规划")
        print("2. problem4_arm_path.png - 显示机械臂路径")
        print("="*50)
        
    except Exception as e:
        logging.error(f"运行出错: {str(e)}")
        import traceback
        traceback.print_exc()
        logging.error(traceback.format_exc())

if __name__ == "__main__":
    main() 